﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net;
using System.Net.Sockets;
using System.Threading;

using LostLib;
using LostLib.Net;


namespace lportmap
{
	class lportmap : IDisposable
	{
		#region Program
		static IPAddress localIP;
		static int localPort;
		static IPAddress remoteIP;
		static int remotePort;
		static byte[] key;

		static void Parse(List<string> args)
		{
			IEnumerator<string> enumerator = args.GetEnumerator();
			while (enumerator.MoveNext())
			{
				switch (enumerator.Current)
				{
				case "-ip":
					if (localIP != null) throw new ArgumentException("duplicate -ip");
					localIP = IPAddress.Parse(enumerator.GetNext());
					break;
				case "-remoteIP":
					if (remoteIP != null) throw new ArgumentException("duplicate -remoteIP");
					remoteIP = IPAddress.Parse(enumerator.GetNext());
					break;
				case "-remotePort":
					if (remotePort != 0) throw new ArgumentException("duplicate -remotePort");
					remotePort = int.Parse(enumerator.GetNext());
					break;
				case "-key":
					if (key != null) throw new ArgumentException("duplicate -key");
					key = Encoding.Default.GetBytes(enumerator.GetNext());
					break;
				case "-port":
					if (localPort != 0) throw new ArgumentException("duplicate -port");
					localPort = int.Parse(enumerator.GetNext());
					break;
				default:
					throw new ArgumentException("Unknown argument: " + enumerator.Current);
				}
			}

			if (key == null)
			{
				key = new byte[16];
				key.Initialize();
			}

			if (localIP == null)
			{
				var entry = Dns.GetHostEntry("localhost");
				localIP = entry.AddressList.First();
			}
			if (remoteIP == null)
				throw new ArgumentNullException("-remoteIP");
			if (remotePort == 0)
				throw new ArgumentNullException("-remotePort");
			if (localPort == 0) localPort = 666;
		}

		static void Main(string[] args)
		{
			Parse(new List<string>(args));

			var listener = new TcpListener(localIP, localPort);

			listener.Start();
			Console.WriteLine("running on {0}:{1}...", localIP, localPort);

			while (true)
			{
				var tcpclient = listener.AcceptTcpClient();
				try
				{
					var client = new lportmap(tcpclient);
					client.TransferError += new EventHandler<ExceptionHandlerEventArgs>(client_TransferError);
					client.RemoteIP = remoteIP;
					client.RemotePort = remotePort;
					client.Key = key;

					Console.WriteLine("{0} connected", tcpclient.Client.RemoteEndPoint);

					ThreadPool.QueueUserWorkItem(new WaitCallback(AcceptClient), client);
				} catch { }

			}
		}

		static void client_TransferError(object sender, ExceptionHandlerEventArgs e)
		{
			var transferError = e.Exception as NetworkTransferException;
			if (transferError != null)
			{
				Console.WriteLine("{0} {1}", transferError.ServerToClient ? "-->" : "<--", transferError.Message);
				e.Reaction = ExceptionReaction.Abort;
				return;
			}
			Console.WriteLine("{0}: {1}", e.Exception.GetType(), e.Exception.Message);
		}

		static void AcceptClient(object obj)
		{
			var client = obj as lportmap;
			client.Run();
		}
		#endregion Program

		private lportmap(TcpClient client)
		{
			Client = client;
		}

		public TcpClient Client { get; private set; }
		public byte[] Key { get; set; }
		public IPAddress RemoteIP { get; set; }
		public int RemotePort { get; set; }

		const byte Zero = 0;
		const int BufferSize = 64 * 1024;

		TcpClient _remote;
		NetworkStream _stream;
		int inKey, outKey;

		private DBNull EncodeIn(byte[] data, int pos, int count)
		{
			while (count > 0)
			{
				data[pos] ^= key[inKey];
				pos++;
				count--;
				inKey++;
				if (inKey >= key.Length) inKey = 0;
			}
			return null;
		}

		private DBNull EncodeOut(byte[] data, int pos, int count)
		{
			while (count > 0)
			{
				data[pos] ^= key[outKey];
				pos++;
				count--;
				outKey++;
				if (outKey >= key.Length) outKey = 0;
			}
			return null;
		}

		public void Run()
		{
			using (_remote = new TcpClient())
			using (_stream = Client.GetStream())
			{
				Client.Client.NoDelay = true;
				try
				{
					_remote.Connect(RemoteIP, RemotePort);
				} catch (SocketException error)
				{
					OnTransferError(new ExceptionHandlerEventArgs(error));
					return;
				}

				ThreadPool.QueueUserWorkItem(new WaitCallback(ServerToClient));

				try
				{
					var other = _remote.GetStream();
					_remote.Client.NoDelay = true;
					Client.NetworkTransfer(_stream, other, BufferSize, EncodeIn);
					Client.Close();
					_remote.Close();
				} catch (Exception error)
				{
					var exception = new NetworkTransferException(error.Message, false);
					var args = new ExceptionHandlerEventArgs(exception);
					OnTransferError(args);
					Client.Close();
					_remote.Close();
					switch (args.Reaction)
					{
					case ExceptionReaction.Ignore:
					case ExceptionReaction.Abort:
						return;
					default:
						throw error;
					}
				}
			}
		}

		private void ServerToClient(object arg)
		{
			try
			{
				using (var remoteStream = _remote.GetStream())
				{
					_remote.NetworkTransfer(remoteStream, _stream, BufferSize, EncodeOut);
					Client.Close();
					_remote.Close();
				}
			} catch (Exception error)
			{
				var exception = new NetworkTransferException(error.Message, false);
				var args = new ExceptionHandlerEventArgs(exception);
				OnTransferError(args);
				_remote.Close();
				Client.Close();
				switch (args.Reaction)
				{
				case ExceptionReaction.Ignore:
				case ExceptionReaction.Abort:
					return;
				default:
					throw error;
				}
			}
		}

		#region Error Handling
		public event EventHandler<ExceptionHandlerEventArgs> TransferError;
		protected virtual void OnTransferError(ExceptionHandlerEventArgs args)
		{
			if (TransferError != null)
				TransferError(this, args);
		}
		#endregion Error Handling

		#region IDisposable Members

		public void Dispose()
		{
			Dispose(true);
			// This object will be cleaned up by the Dispose method.
			// Therefore, you should call GC.SupressFinalize to
			// take this object off the finalization queue
			// and prevent finalization code for this object
			// from executing a second time.
			GC.SuppressFinalize(this);
		}

		// Dispose(bool disposing) executes in two distinct scenarios.
		// If disposing equals true, the method has been called directly
		// or indirectly by a user's code. Managed and unmanaged resources
		// can be disposed.
		// If disposing equals false, the method has been called by the
		// runtime from inside the finalizer and you should not reference
		// other objects. Only unmanaged resources can be disposed.
		private void Dispose(bool disposing)
		{
			// Check to see if Dispose has already been called.
			if (Client != null)
			{
				// If disposing equals true, dispose all managed
				// and unmanaged resources.
				if (disposing)
				{
					// Dispose managed resources.
					Client.Close();
				}

				// Call the appropriate methods to clean up
				// unmanaged resources here.
				// If disposing is false,
				// only the following code is executed.

				Client = null;
			}
		}

		// Use C# destructor syntax for finalization code.
		// This destructor will run only if the Dispose method
		// does not get called.
		// It gives your base class the opportunity to finalize.
		// Do not provide destructors in types derived from this class.
		~lportmap()
		{
			// Do not re-create Dispose clean-up code here.
			// Calling Dispose(false) is optimal in terms of
			// readability and maintainability.
			Dispose(false);
		}


		#endregion
	}
}
